from concurrent.futures import ThreadPoolExecutor
import cv2
import numpy as np
import time
from .utils import calculate_similarity
from tqdm import tqdm
import os

class ImageDeduplicator:
    """图片去重器
        
    这个类用于处理视频帧或图片序列的去重，主要功能包括：
    1. 基于SIFT特征的图片相似度计算
    2. 自动场景分组和质量评估
    3. 智能选择最佳质量帧
    
    算法优点：
    - 特征提取稳定：SIFT特征对旋转、缩放、亮度变化具有良好的不变性
    - 场景识别准确：通过连续多帧相似度判断，避免单帧误判
    - 质量评估全面：综合考虑清晰度、对比度和亮度均匀性
    - 内存占用优化：及时释放不需要的帧数据
    - 输入格式灵活：支持多种输入格式
    
    算法缺点：
    - 计算开销大：SIFT特征提取和匹配计算量较大
    - 内存消耗：处理高分辨率图片时需要较大内存
    - 阈值敏感：相似度阈值的设置会影响分组效果
    - 场景切换延迟：需要连续多帧才能确认场景变化
    
    
    这个类用于处理视频帧或图片序列的去重，主要功能包括：
    1. 基于SIFT特征的图片相似度计算
    2. 自动场景分组和质量评估
    3. 智能选择最佳质量帧
    
    主要特点：
    - 使用SIFT特征进行图片相似度比较
    - 支持自动场景切换检测（连续10帧差异大于阈值）
    - 综合评估图片质量（清晰度、对比度、亮度均匀性）
    - 内存优化，自动释放不需要的帧
    - 支持字典格式和简单路径格式的输入
    
    参数说明：
    - similarity_threshold: 相似度阈值，默认0.6
    - base_window: 场景检测窗口大小，默认10帧
    - batch_size: 批处理大小，默认10
    
    调用流程：
    1. 初始化：
       deduplicator = ImageDeduplicator(similarity_threshold=0.6)
    
    2. 输入帧路径列表：
       frame_paths = [
           {'frame_path': 'path/to/frame1.jpg'},
           {'frame_path': 'path/to/frame2.jpg'},
           ...
       ]
       或者简单的路径列表：
       frame_paths = ['path/to/frame1.jpg', 'path/to/frame2.jpg', ...]
    
    3. 执行去重：
       unique_frames = deduplicator.deduplicate_frames(frame_paths)
    
    去重算法流程：
    1. 以第一帧为基准帧，计算后续帧与基准帧的SIFT特征相似度
    2. 当连续10帧相似度都小于0.4时，认为遇到新的场景
    3. 对每组相似帧，选择质量最高的一帧保留
    4. 使用processed_paths确保每个图片只会被分到一个组中
    
    返回结果：
    - 返回去重后的帧列表，每个元素保持输入的原始格式
    - 如果输入是字典列表，返回字典列表；如果是路径列表，返回路径列表
    """
    def __init__(self, similarity_threshold=0.6):
        self.similarity_threshold = similarity_threshold
        self._sift_feature_cache = {}
        self.max_cache_size = 300
        self.base_window = 10
        self.cache_timeout = 900
        self.batch_size = 10
        
        # 初始化SIFT特征提取器
        self.sift = cv2.SIFT_create()
        # FLANN参数设置
        FLANN_INDEX_KDTREE = 1
        self.index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
        self.search_params = dict(checks=50)
        self.flann = cv2.FlannBasedMatcher(self.index_params, self.search_params)

    def _extract_sift_features(self, img):
        """提取SIFT特征"""
        try:
            if img is None:
                return None, None
                
            # 转换为灰度图
            if len(img.shape) > 2:
                gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            else:
                gray = img
                
            # 降采样处理以提高性能
            scale_factor = 0.5
            gray = cv2.resize(gray, None, fx=scale_factor, fy=scale_factor)
            
            # 提取SIFT关键点和描述符
            keypoints, descriptors = self.sift.detectAndCompute(gray, None)
            
            return keypoints, descriptors
            
        except Exception as e:
            print(f"SIFT特征提取错误: {str(e)}")
            return None, None

    def _match_sift_features(self, desc1, desc2):
        """匹配SIFT特征"""
        try:
            if desc1 is None or desc2 is None or len(desc1) < 2 or len(desc2) < 2:
                return 0.0
                
            # 使用FLANN匹配器进行特征匹配
            matches = self.flann.knnMatch(desc1, desc2, k=2)
            
            # 应用比率测试筛选良好匹配
            good_matches = []
            for m, n in matches:
                if m.distance < 0.75 * n.distance:
                    good_matches.append(m)
            
            # 计算相似度分数
            similarity = len(good_matches) / max(len(desc1), len(desc2))
            
            return min(similarity * 3.0, 1.0)  # 放大相似度并限制在[0,1]范围内
            
        except Exception as e:
            print(f"SIFT特征匹配错误: {str(e)}")
            return 0.0

    def _calculate_similarity_score(self, img1, img2):
        """计算两帧之间的相似度"""
        try:
            # 提取SIFT特征
            _, desc1 = self._extract_sift_features(img1)
            _, desc2 = self._extract_sift_features(img2)
            
            # 计算SIFT特征相似度
            similarity = self._match_sift_features(desc1, desc2)
            
            return similarity
            
        except Exception as e:
            print(f"相似度计算错误: {str(e)}")
            return 0.0

    def deduplicate_frames(self, frame_paths):
        """基于SIFT特征和质量的分组去重"""
        try:
            if not frame_paths:
                return []
            
            unique_frames = []
            processed_paths = set()  # 记录已处理的图片路径
            i = 0
            
            while i < len(frame_paths):
                # 获取基准帧
                base_path = frame_paths[i]['frame_path'] if isinstance(frame_paths[i], dict) else str(frame_paths[i])
                
                # 如果该图片已经处理过，跳过
                if base_path in processed_paths:
                    i += 1
                    continue
                
                base_frame = cv2.imread(base_path)
                if base_frame is None:
                    i += 1
                    continue
                
                # 存储当前组的帧
                current_group = [frame_paths[i]]
                processed_paths.add(base_path)
                j = i + 1
                low_similarity_count = 0
                
                # 向后遍历直到找到连续多帧差异较大的情况
                last_dissimilar_frame = j
                while j < len(frame_paths):
                    curr_path = frame_paths[j]['frame_path'] if isinstance(frame_paths[j], dict) else str(frame_paths[j])
                    
                    # 如果该图片已经处理过，跳过
                    if curr_path in processed_paths:
                        j += 1
                        continue
                        
                    curr_frame = cv2.imread(curr_path)
                    if curr_frame is None:
                        j += 1
                        continue
                    
                    similarity = self._calculate_similarity_score(base_frame, curr_frame)
                    print(f"基准帧 {i} -> 当前帧 {j} 相似度: {similarity:.3f}")
                    
                    if similarity < 0.4:
                        last_dissimilar_frame = j
                        low_similarity_count += 1
                        if low_similarity_count >= 5:
                            break
                    else:
                        low_similarity_count = 0
                        current_group.append(frame_paths[j])
                        processed_paths.add(curr_path)  # 将相似帧标记为已处理
                    
                    j += 1
                    del curr_frame
                
                # 从当前组中选择质量最高的帧
                best_frame = self._select_best_quality_frame(current_group)
                if best_frame:
                    unique_frames.append(best_frame)
                
                # 从最后一个不相似帧开始处理
                i = last_dissimilar_frame
                del base_frame
            
            return unique_frames
            
        except Exception as e:
            print(f"去重帧序列错误: {str(e)}")
            return []

    def _select_best_quality_frame(self, frame_group):
        """从一组相似帧中选择质量最高的"""
        try:
            best_quality = -1
            best_frame = None
            
            for frame_info in frame_group:
                frame_path = frame_info['frame_path'] if isinstance(frame_info, dict) else str(frame_info)
                frame = cv2.imread(frame_path)
                if frame is None:
                    continue
                
                # 计算帧质量
                quality = self._calculate_frame_quality(frame)
                print(f"帧 {os.path.basename(frame_path)} 质量: {quality:.3f}")
                
                if quality > best_quality:
                    best_quality = quality
                    best_frame = frame_info
                
                del frame
            
            return best_frame
            
        except Exception as e:
            print(f"选择最佳帧错误: {str(e)}")
            return None

    def analyze_frame_content(self, frame):
        """分析帧内容"""
        try:
            if frame is None:
                return {
                    'quality': 0.0,
                    'doc_detected': False
                }
            
            # 计算帧质量
            quality = self._calculate_frame_quality(frame)
            
            return {
                'quality': quality,
                'doc_detected': quality > 0.2
            }
            
        except Exception as e:
            print(f"分析帧内容错误: {str(e)}")
            return {
                'quality': 0.0,
                'doc_detected': False
            }

    def _calculate_frame_quality(self, frame):
        """计算帧质量"""
        try:
            if frame is None:
                return 0.0
            
            # 转换为灰度图
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            
            # 计算清晰度（拉普拉斯方差）
            laplacian = cv2.Laplacian(gray, cv2.CV_64F)
            sharpness = np.var(laplacian) / 10000  # 归一化
            
            # 计算对比度
            contrast = np.std(gray) / 128
            
            # 计算亮度均匀性
            brightness = np.mean(gray) / 255
            brightness_uniformity = 1 - np.std(gray) / 128
            
            # 综合评分
            quality_score = (
                sharpness * 0.4 +
                contrast * 0.3 +
                brightness_uniformity * 0.3
            )
            
            return min(max(quality_score, 0.0), 1.0)
            
        except Exception as e:
            print(f"计算帧质量错误: {str(e)}")
            return 0.0
